# predict/predict_onnx.py
"""
Python推理端（ONNX Runtime）
- 加载训练好的LightGBM ONNX模型
- 使用与训练相同的特征标准化
- 对每只股票预测次日开盘上涨概率与收盘上涨概率
- 结果写入PostgreSQL数据库
"""

import pandas as pd                     # 数据处理库
import numpy as np                      # 数值计算库
import joblib                           # 模型序列化库
import onnxruntime as ort               # ONNX运行时库
from sqlalchemy import create_engine, text  # 数据库连接库
from config import DB_CONFIG            # 数据库配置

# 特征列，需与训练端一致
FEATURE_COLS = [
    "ret_1d", "ret_5d",                 # 收益率特征
    "ma5", "ma10",                      # 移动平均线特征
    "rsi14",                            # 相对强弱指数
    "vol_ma5",                          # 成交量移动平均
    "sse_ret", "csi300_ret", "industry_ret"  # 市场和行业基准收益
]

# 模型与Scaler路径配置
SCALER_PATH = "scaler.pkl"              # 特征标准化器路径
ONNX_OPEN_PATH = "model_open.onnx"      # 开盘预测模型路径
ONNX_CLOSE_PATH = "model_close.onnx"    # 收盘预测模型路径
MODEL_VERSION = "v1.0"                  # 模型版本号

# 1. 加载特征标准化对象（用于将特征数据标准化到训练时的分布）
scaler = joblib.load(SCALER_PATH)

# 2. 初始化ONNX Runtime Session（创建模型推理会话）
sess_open = ort.InferenceSession(ONNX_OPEN_PATH)    # 开盘预测模型会话
sess_close = ort.InferenceSession(ONNX_CLOSE_PATH)  # 收盘预测模型会话

# 3. 创建数据库引擎（用于连接PostgreSQL数据库）
engine = create_engine(
    f"postgresql+psycopg2://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
)

# 4. 读取待推理特征（从features_daily表中获取所有特征数据）
with engine.connect() as conn:
    df = pd.read_sql(text("SELECT * FROM features_daily ORDER BY trade_date"), conn)

# 5. 标准化特征（使用训练时保存的scaler对特征进行标准化处理）
X = scaler.transform(df[FEATURE_COLS].values.astype(np.float32))

# 6. 推理得到上涨概率
# 假设ONNX模型输出是二维概率矩阵：[prob_down, prob_up]
# 模型输出两个类别的概率，我们关注上涨概率（索引为1）
proba_open = sess_open.run(None, {sess_open.get_inputs()[0].name: X})[0]   # 开盘上涨概率
proba_close = sess_close.run(None, {sess_close.get_inputs()[0].name: X})[0]  # 收盘上涨概率

# 7. 转换为整数百分比（将概率转换为0-100的整数百分比）
df["open_up_pct"] = (proba_open[:, 1] * 100).astype(int)    # 开盘上涨概率百分比
df["close_up_pct"] = (proba_close[:, 1] * 100).astype(int)  # 收盘上涨概率百分比

# 8. 准备输出结果（只保留需要的列）
df_out = df[["trade_date", "ticker", "open_up_pct", "close_up_pct"]].copy()

# 9. 写入PostgreSQL predictions表（将预测结果保存到数据库）
with engine.begin() as conn:
    for _, row in df_out.iterrows():
        conn.execute(text("""
            INSERT INTO predictions (ticker, trade_date, predict_time, open_up_pct, close_up_pct, model_version)
            VALUES (:ticker, :trade_date, now(), :open_up_pct, :close_up_pct, :model_version)
            ON CONFLICT (ticker, trade_date) DO UPDATE
              SET open_up_pct = EXCLUDED.open_up_pct,
                  close_up_pct = EXCLUDED.close_up_pct,
                  predict_time = now()
        """), {
            "ticker": row["ticker"],              # 股票代码
            "trade_date": row["trade_date"],      # 交易日期
            "open_up_pct": row["open_up_pct"],    # 开盘上涨概率百分比
            "close_up_pct": row["close_up_pct"],  # 收盘上涨概率百分比
            "model_version": MODEL_VERSION        # 模型版本
        })

print(f"✅ 推理完成，结果已写入 predictions 表，模型版本: {MODEL_VERSION}")